package com.ryan.ohos.extension.state;

import ohos.aafwk.ability.Ability;
import ohos.agp.components.Component;
import ohos.agp.components.ComponentContainer;
import ohos.app.Context;
import ohos.utils.PacMap;

import java.lang.reflect.Field;

public class StateHelper {
    public static void saveState(Ability context, PacMap state) {
        ComponentContainer root = getRoot(context);
        if (root == null) return;

        callState(root, state, true);
    }

    public static void restoreState(Ability context, PacMap state) {
        ComponentContainer root = getRoot(context);
        if (root == null) return;

        callState(root, state, false);
    }

    private static void callState(ComponentContainer container, PacMap state, boolean save) {
        int count = container.getChildCount();
        for (int i = 0; i < count; i++) {
            Component child = container.getComponentAt(i);
            if (child instanceof SaveRestoreState) {
                SaveRestoreState saveRestoreState = (SaveRestoreState) child;
                if (save) {
                    saveRestoreState.onSaveState(state);
                } else {
                    saveRestoreState.onRestoreState(state);
                }
            }

            if (child instanceof ComponentContainer) {
                callState((ComponentContainer) child, state, save);
            }
        }
    }

    private static ComponentContainer getRoot(Ability context) {
        try {
            Class<?> clz = context.getWindow().getClass();

            Field declaredField = clz.getDeclaredField("agpWindow");
            declaredField.setAccessible(true);
            Object window = declaredField.get(context.getWindow());

            Class<?> windowClz = window.getClass();
            Field declaredField1 = windowClz.getDeclaredField("mWindowRoot");
            declaredField1.setAccessible(true);

            return (ComponentContainer) declaredField1.get(window);
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }
}
